package com.ideaaedi.hot.clazz.update;

import com.ideaaedi.commonds.clazz.ClassNameGrabber;
import com.ideaaedi.commonds.io.IOUtil;
import com.ideaaedi.component.compile.Compiler;
import com.ideaaedi.component.decompile.Decompiler;
import com.ideaaedi.hot.clazz.update.enums.TransferTypeEnum;
import com.ideaaedi.hot.clazz.update.exception.HotUpdateClassException;
import com.ideaaedi.hot.clazz.update.memory.MemoryTransferMap;
import com.ideaaedi.hot.common.HotUpdateConstant;
import org.apache.commons.lang3.StringUtils;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.File;
import java.io.PrintWriter;
import java.io.StringWriter;
import java.lang.instrument.ClassDefinition;
import java.lang.instrument.Instrumentation;
import java.lang.instrument.UnmodifiableClassException;
import java.lang.reflect.UndeclaredThrowableException;
import java.nio.charset.StandardCharsets;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Base64;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.stream.Collectors;

/**
 * class热更新代理
 * <p>
 * <p>
 * Java Instrumentation可以获取到JVM中的类信息，进而作出修改，这样就实现了热更新。Instrumentation操作修改字节码的方式有两种：
 * <ul>
 *     <li>pre-main：在程序执行main方法之前，可以修改字节码。pre-main的关键方法如下，你可以任选一个实现，
 *         如果都实现，那么会优先执行带有Instrumentation参数的方法，然后再执行不带该参数的方法
 *         <ol>
 *             <li><code>public static void premain(String agentArgs, Instrumentation inst)</code></li>
 *             <li><code>public static void premain(String agentArgs)</code></li>
 *         </ol>
 *     </li>
 *     <li>agent-main：在程序执行main方法之后，可以修改字节码。agent-main的关键方法如下，你可以任选一个实现，也可以都实现，
 *         如果都实现，那么会优先执行带有Instrumentation参数的方法，然后再执行不带该参数的方法
 *         <ol>
 *             <li><code>public static void agentmain (String agentArgs, Instrumentation inst)</code></li>
 *             <li><code>public static void agentmain (String agentArgs)</code></li>
 *         </ol>
 *     </li>
 * </ul>
 *
 * <p>
 * Instrumentation的局限性：
 * <ul>
 *     <li>premain和agentmain两种方式修改字节码的时机都是类文件加载之后，也就是说必须要带有Class类型的参数，不能通过字节码文件和自定义的类名重新定义一个本来不存在的类。</li>
 *     <li>类的字节码修改称为类转换(Class Transform)，类转换其实最终都回归到类重定义Instrumentation#redefineClasses()方法，此方法有以下限制：
 *         <ol>
 *             <li>新类和老类的父类必须相同</li>
 *             <li>新类和老类实现的接口数也要相同，并且是相同的接口</li>
 *             <li>新类和老类访问符必须一致。 新类和老类字段数和字段名要一致</li>
 *             <li>新类和老类新增或删除的方法必须是private static/final修饰的</li>
 *             <li>可以修改方法体</li>
 *         </ol>
 *     </li>
 * </ul>
 *
 * <p>
 * 更多Instrumentation的介绍，可见好文
 * <a href="https://www.cnblogs.com/rickiyang/p/11368932.html">https://www.cnblogs.com/rickiyang/p/11368932.html</a>
 *
 * @author JustryDeng
 * @since 1.0.0
 */
public class ClassHotUpdateAgent {
    
    public static final Logger log = LoggerFactory.getLogger(ClassHotUpdateAgent.class);
    
    /**
     * agent-main
     */
    public static void agentmain(String agentArgs, Instrumentation ins) {
        try {
            log.info("Step into ClassHotUpdateAgent#agentmain, agentArgs -> {}", agentArgs);
            long startTime = System.currentTimeMillis();
            List<byte[]> classBytes = obtainClassBytes(agentArgs);
            long obtainClassBytesCompletedTime = System.currentTimeMillis();
            log.info("Complete obtainClassBytes. consume time -> {}ms.", obtainClassBytesCompletedTime - startTime);
            hotUpdate(classBytes, ins);
            long endTime = System.currentTimeMillis();
            log.info("Complete hotUpdate. consume time -> {}ms.", endTime - startTime);
        } catch (Exception e) {
            e.printStackTrace();
            String errMsg = "hotUpdate occur exception " + obtainStackTraceMessage(e);
            log.error(errMsg);
        }
    }
    
    public static void hotUpdate(List<byte[]> classBytes, Instrumentation ins) {
        // key-全类名, value-class字节流
        Map<String, byte[]> classMap = new HashMap<>(8);
        classBytes.forEach(classByte -> classMap.put(
                ClassNameGrabber.obtainFromSourceCode(Decompiler.defaultDecompiler().decompile(classByte)),
                classByte));
        Set<String> toUpdateClassNameSet = classMap.keySet();
        log.info("Start hotUpdate, classFilePaths -> {}", toUpdateClassNameSet);
        
        //noinspection rawtypes
        Class[] allLoadedClasses = ins.getAllLoadedClasses();
        List<String> allLoadedClassNameList = Arrays.stream(allLoadedClasses)
                .map(Class::getName).sorted().collect(Collectors.toList());
        log.debug("allLoadedClassNameList -> {}", allLoadedClassNameList);
        try {
            //noinspection rawtypes
            for (Class alreadyLoadedClass : allLoadedClasses) {
                String className = alreadyLoadedClass.getName();
                if (toUpdateClassNameSet.contains(className)) {
                    ClassDefinition classDefinition = new ClassDefinition(alreadyLoadedClass, classMap.get(className));
                    // 执行替换
                    ins.redefineClasses(classDefinition);
                    log.info("redefine class success for {}", className);
                }
            }
        } catch (ClassNotFoundException | UnmodifiableClassException e) {
            throw new UndeclaredThrowableException(e);
        }
    }
    
    /**
     * 获取到.class字节数组
     *
     * @param agentArgs
     *         代理入参
     *
     * @return .class字节数组集合
     */
    public static List<byte[]> obtainClassBytes(String agentArgs) {
        log.info("Start obtainClassBytes agentArgs -> {}", agentArgs);
        if (StringUtils.isBlank(agentArgs)) {
            throw new ParseClassBytesException("agentArgs cannot be blank.");
        }
        if (agentArgs.startsWith(TransferTypeEnum.MEMORY_JAVA_BYTES.getProtocol())) {
            log.info("hit -> {}", TransferTypeEnum.MEMORY_JAVA_BYTES.getProtocol());
            String javaByteKeys = agentArgs.substring(TransferTypeEnum.MEMORY_JAVA_BYTES.getProtocol().length());
            if (StringUtils.isBlank(javaByteKeys)) {
                throw new ParseClassBytesException("javaByteKeys cannot be blank.");
            }
            // 获取到java源代码
            List<byte[]> sourceCodes = Arrays.stream(javaByteKeys.split(HotUpdateConstant.COMMA))
                    .filter(StringUtils::isNotBlank)
                    .map(base64Key -> new String(
                                    Base64.getDecoder().decode(base64Key.getBytes(StandardCharsets.UTF_8)),
                                    StandardCharsets.UTF_8
                          )
                    )
                    .map(javaByteKey -> {
                        byte[] bytes = MemoryTransferMap.globalMemory().get(javaByteKey);
                        if (bytes == null) {
                            throw new HotUpdateClassException(String.format("bytes is null, curr key is -> %s, all key is %s.",
                                    javaByteKey, MemoryTransferMap.globalMemory().keySet()));
                        }
                        return bytes;
                    })
                    .collect(Collectors.toList());
            // 动态编译源码为.class的byte[]形式
            Compiler compiler = Compiler.create(Thread.currentThread().getContextClassLoader());
            for (byte[] javaSourceCode : sourceCodes) {
                compiler.addSource(javaSourceCode);
            }
            return new ArrayList<>(compiler.buildByteCodes().values());
        }
        if (agentArgs.startsWith(TransferTypeEnum.BASE64_JAVA_PATHS.getProtocol())) {
            log.info("hit -> {}", TransferTypeEnum.BASE64_JAVA_PATHS.getProtocol());
            String base64JavaPaths = agentArgs.substring(TransferTypeEnum.BASE64_JAVA_PATHS.getProtocol().length());
            if (StringUtils.isBlank(base64JavaPaths)) {
                throw new ParseClassBytesException("base64JavaPaths cannot be blank.");
            }
            // 获取到java源代码
            List<String> sourceCodeFilePath = Arrays.stream(base64JavaPaths.split(HotUpdateConstant.COMMA))
                    .filter(StringUtils::isNotBlank)
                    .map(base64JavaPath -> new String(
                                    Base64.getDecoder().decode(base64JavaPath.getBytes(StandardCharsets.UTF_8)),
                                    StandardCharsets.UTF_8
                            )
                    ).collect(Collectors.toList());
            log.info("sourceCodeFilePath -> {}", sourceCodeFilePath);
            List<byte[]> sourceCodes = sourceCodeFilePath.stream()
                    .map(path -> IOUtil.toBytes(new File(path)))
                    .collect(Collectors.toList());
            // 动态编译源码为.class的byte[]形式
            Compiler compiler = Compiler.create(Thread.currentThread().getContextClassLoader());
            for (byte[] javaSourceCode : sourceCodes) {
                compiler.addSource(javaSourceCode);
            }
            return new ArrayList<>(compiler.buildByteCodes().values());
        }
        if (agentArgs.startsWith(TransferTypeEnum.MEMORY_CLASS_BYTES.getProtocol())) {
            log.info("hit -> {}", TransferTypeEnum.MEMORY_CLASS_BYTES.getProtocol());
            String classByteKeys = agentArgs.substring(TransferTypeEnum.MEMORY_CLASS_BYTES.getProtocol().length());
            if (StringUtils.isBlank(classByteKeys)) {
                throw new ParseClassBytesException("classByteKeys cannot be blank.");
            }
            return Arrays.stream(classByteKeys.split(HotUpdateConstant.COMMA))
                    .filter(StringUtils::isNotBlank)
                    .map(base64Key -> new String(
                                    Base64.getDecoder().decode(base64Key.getBytes(StandardCharsets.UTF_8)),
                                    StandardCharsets.UTF_8
                            )
                    )
                    .map(classByteKey -> {
                        byte[] bytes = MemoryTransferMap.globalMemory().get(classByteKey);
                        if (bytes == null) {
                            throw new HotUpdateClassException(String.format("bytes is null, curr key is -> %s, all key is %s.",
                                    classByteKey, MemoryTransferMap.globalMemory().keySet()));
                        }
                        return bytes;
                    })
                    .collect(Collectors.toList());
        }
        if (agentArgs.startsWith(TransferTypeEnum.BASE64_CLASS_PATHS.getProtocol())) {
            log.info("hit -> {}", TransferTypeEnum.BASE64_CLASS_PATHS.getProtocol());
            String base64ClassPaths = agentArgs.substring(TransferTypeEnum.BASE64_CLASS_PATHS.getProtocol().length());
            if (StringUtils.isBlank(base64ClassPaths)) {
                throw new ParseClassBytesException("base64ClassPaths cannot be blank.");
            }
            return Arrays.stream(base64ClassPaths.split(HotUpdateConstant.COMMA))
                    .filter(StringUtils::isNotBlank)
                    .map(base64ClassPath ->
                            new String(
                                    Base64.getDecoder().decode(base64ClassPath.getBytes(StandardCharsets.UTF_8)),
                                    StandardCharsets.UTF_8
                            )
                    ).map(base64ClassPath -> IOUtil.toBytes(new File(base64ClassPath)))
                    .collect(Collectors.toList());
        }
        // 不是以指定协议开头的，则抛出异常
        throw new ParseClassBytesException(String.format("agentArgs must startWith protocol '%s' or '%s' or '%s' or "
                        + "'%s'.",
                TransferTypeEnum.MEMORY_JAVA_BYTES.getProtocol(),
                TransferTypeEnum.BASE64_JAVA_PATHS.getProtocol(),
                TransferTypeEnum.MEMORY_CLASS_BYTES.getProtocol(),
                TransferTypeEnum.BASE64_CLASS_PATHS.getProtocol()
        ));
    }
    
    /**
     * 将异常堆栈 信息 转换为字符串
     *
     * @param e
     *         异常
     *
     * @return 该异常的错误堆栈信息
     */
    public static String obtainStackTraceMessage(Exception e) {
        try (StringWriter sw = new StringWriter(); PrintWriter pw = new PrintWriter(sw)) {
            // 将异常的的堆栈信息输出到printWriter中
            e.printStackTrace(pw);
            pw.flush();
            sw.flush();
            return sw.toString();
        } catch (Exception exception) {
            // ignore
            e.printStackTrace();
            log.error("hotUpdate occur exception", e);
            return "getStackTraceMessage occur exception" + exception.getMessage();
        }
    }
}
